#
# plot.py
#
# Copyright (C) 2012 Robert Buj Gelonch
# Copyright (C) 2012 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__author__ = "Robert Buj Gelonch, and David Megias Jimenez"
__copyright__ = "Copyright 2012, Robert Buj Gelonch and David Megias Jimenez"
__credits__ = ["Robert Buj Gelonch", "David Megias Jimenez"]
__license__ = "GPL"
__version__ = "3"
__maintainer__ = "Robert Buj"
__email__ = "rbuj@uoc.edu"
__status__ = "Development"
__docformat__ = 'plaintext'

import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
import numpy
from pylab import plot
from pylab import show
from pylab import subplot
from pylab import title
from pylab import xlabel
from pylab import ylabel
from scipy import arange
from scipy import fft
from scipy.signal import freqz

def signal_spectrum(sample_rate, samples, fig):
    n = len(samples) # length of the signal
    k = arange(n) # samples index vector
    Ts = 1.0 / sample_rate
    T = n / sample_rate
    t = arange(0, n * Ts, Ts) # time vector

    #--------------------------------------------------------------
    # Frequency response
    #--------------------------------------------------------------
    frq = k / T # two sides frequency range
    frq = frq[range(n / 2)] # one side frequency range
    Y = fft(samples) / n # fft computing and normalization
    Y = Y[range(n / 2)]

    #--------------------------------------------------------------
    # Plotting
    #--------------------------------------------------------------
    # Time plot
    ax = fig.add_subplot(211)
    ax.set_title('Signal')
    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Amplitude')
    ax.plot(t, samples, 'b') # plotting the signal

    # Freq plot
    subplot(212)
    title('Frequency spectrum')
    xlabel('Freq (Hz)')
    ylabel('|Y(freq)|')
    plot(frq, abs(Y), 'r') # plotting the spectrum
    show()

def filter(filtered_signal, samples, sample_rate, numtaps, fir_coeff, nyq_rate, title):
    n = len(samples) # length of the signal
    k = arange(n) # samples index vector
    Ts = 1.0 / sample_rate
    T = n / sample_rate
    t = arange(0, n * Ts, Ts) # time vector

    frq = k / T # two sides frequency range
    frq = frq[range(n / 2)] # one side frequency range
    Y = fft(samples) / n # fft computing and normalization
    Y = Y[range(n / 2)]

    #--------------------------------------------------------------
    # Frequency response of the signal filtered
    #--------------------------------------------------------------
    Y2 = fft(filtered_signal) / n # fft computing and normalization
    Y2 = Y2[range(n / 2)]

    #--------------------------------------------------------------
    # Plotting
    #--------------------------------------------------------------
    fig = plt.figure(num=1, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
    fig.suptitle(title)
    # Time plot
    plt.subplot(221)
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.title('Signal')
    plt.plot(t, samples, 'b') # plotting the signal

    # Freq plot
    plt.subplot(222)
    plt.xlabel('Freq (Hz)')
    plt.ylabel('|Y(freq)|')
    plt.title('Frequency Spectrum')
    plt.plot(frq, abs(Y), 'b') # plotting the spectrum

    # Filtered signal
    plt.subplot(221)
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.plot(t, filtered_signal, 'r') # plotting the signal

    # Freq plot
    plt.subplot(222)
    plt.xlabel('Freq (Hz)')
    plt.ylabel('|Y(freq)|')
    plt.plot(frq, abs(Y2), 'r') # plotting the spectrum

    # Plot the FIR filter coefficients.
    plt.subplot(223)
    plt.title('Filter Coefficients (%d taps)' % numtaps)
    plt.plot(fir_coeff, 'bo-') # plotting the spectrum

    # Plot the magnitude response of the filter.
    plt.subplot(224)
    w, h = freqz(fir_coeff, worN=8000)
    plt.title('Frequency Response')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Gain')
    plt.ylim(-0.05, 1.05)
    plt.plot((w / numpy.pi) * nyq_rate, abs(h), linewidth=2)
    return plt

def two_signals(cA, marked_message, title, title_a, title_b):
    fig = plt.figure(num=1, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')
    fig.suptitle(title)

    # Plot signal A
    ax = plt.subplot(211)
    plt.ylabel('Amplitude')
    plt.title(title_a)
    plt.plot(cA, 'b')
    ax.set_xticklabels([])
    ax.set_xlim([0, len(cA)])
    plt.grid(True)

    # Plot signal B
    ax = plt.subplot(212)
    plt.ylabel('Amplitude')
    plt.xlabel('Samples')
    plt.title(title_b)
    plt.plot(marked_message, 'b')
    ax.set_xlim([0, len(marked_message)])
    plt.grid(True)
    return plt

def p1_bench_logarithmic_quantification(cA, marked_message, message_secret_val, steps, frame_size, message_secret, secret):
    fig = plt.figure(num=2, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
    fig.suptitle('Adaptive Speech watermarking in wavelet domain based on Logarithm')

    ax = plt.subplot(311)
    ax.xaxis.set_major_locator(MultipleLocator(frame_size))
    ax.set_xticklabels([])
    plt.ylabel('Amplitude')
    plt.title('Marked Frames')
    plt.grid(True)
    plt.stem(range(frame_size * len(secret)), marked_message[:(frame_size * len(secret))], 'b')

    ax = plt.subplot(312)
    ax.xaxis.set_major_locator(MultipleLocator(frame_size))
    ax.set_xticklabels([])
    plt.title('Steganography')
    plt.ylabel('Base')
    plt.grid(True, which='both')
    plt.bar(range(frame_size * len(secret)), message_secret[:(frame_size * len(secret))], 1, color='r', label="decimals")
    plt.bar(range(frame_size * len(secret)), numpy.floor(numpy.log10(abs(cA[:(frame_size * len(secret))])) / numpy.log10(steps)), 1, color='b', label="floor")
    leg = ax.legend(loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)

    ax = plt.subplot(313)
    plt.title('Secret Message')
    plt.ylabel('Value')
    plt.xlabel('Frames')
    ax.set_ylim([0, 1])
    ax.set_autoscaley_on(False)
    ax.set_yticks([0, 1])
    ax.xaxis.set_major_locator(MultipleLocator(frame_size))
    ax.set_xticklabels(range((frame_size * len(secret))))
    plt.grid(True)
    plt.plot(range(frame_size * len(secret)), message_secret_val[:(frame_size * len(secret))], 'bo--')

    return plt
